from sentence_transformers import SentenceTransformer
import torch
from args import *
import time
# Load or create a SentenceTransformer model.
model = SentenceTransformer(MODEL_NAME, cache_folder=MODEL_PATH)
# Get device like 'cuda'/'cpu' that should be used for computation.
if torch.cuda.is_available():
    model = model.to(torch.device("cuda"))
print(f"在{model.device}上运行pytorch")

def encode(documents: list):
    start = time.time()
    print(documents)
    embeddings = model.encode(documents, device='cuda' if torch.cuda.is_available() else 'cpu', batch_size= BATCH_SIZE_MODEL )
    cost = time.time() - start
    print(f"转换{len(documents)}个向量耗时 {cost:.2f}s")
    embeddings_float = embeddings.tolist()
    torch.cuda.empty_cache()
    return embeddings_float
